from PyQt6.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
    QLabel, QFileDialog, QMessageBox, QGroupBox, QProgressBar, QTextEdit
)
from PyQt6.QtGui import QPixmap, QImage, QFont
from PyQt6.QtCore import Qt

class MedicalDetectionUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("医学影像病灶检测系统")
        self.setGeometry(100, 100, 1300, 800)

        # 文件路径
        self.current_image_path = None
        self.detection_results = None
        self.original_pixmap = None
        self.worker = None  # 保存worker引用
        
        # 显示模式：True显示检测图，False显示原图
        self.show_annotated = True

        # 初始化界面
        self.init_ui()

    def init_ui(self):
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QHBoxLayout(main_widget)

        # 左侧控制面板
        left_panel = QGroupBox("控制面板")
        left_layout = QVBoxLayout(left_panel)

        self.upload_btn = QPushButton("上传医学影像")
        self.upload_btn.clicked.connect(self.upload_image)
        left_layout.addWidget(self.upload_btn)

        self.detect_btn = QPushButton("病灶检测")
        self.detect_btn.setEnabled(False)
        left_layout.addWidget(self.detect_btn)
        
        # 添加切换视图按钮
        self.toggle_view_btn = QPushButton("切换视图 (原图/检测图)")
        self.toggle_view_btn.clicked.connect(self.toggle_view)
        self.toggle_view_btn.setEnabled(False)
        left_layout.addWidget(self.toggle_view_btn)

        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        left_layout.addWidget(self.progress_bar)

        self.results_label = QLabel("检测结果:")
        self.results_label.setFont(QFont("Arial", 12, QFont.Weight.Bold))
        left_layout.addWidget(self.results_label)

        self.results_text = QTextEdit()
        self.results_text.setReadOnly(True)
        self.results_text.setMaximumHeight(150)
        left_layout.addWidget(self.results_text)

        left_layout.addStretch()

        # 右侧显示面板
        right_panel = QGroupBox("影像显示")
        right_layout = QVBoxLayout(right_panel)

        self.image_label = QLabel("请上传医学影像")
        self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.image_label.setMinimumSize(600, 500)
        self.image_label.setStyleSheet("border: 1px solid gray;")
        right_layout.addWidget(self.image_label)

        main_layout.addWidget(left_panel, 1)
        main_layout.addWidget(right_panel, 3)

    def upload_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择医学影像",
            "",
            "Image Files (*.png *.jpg *.jpeg *.bmp *.tif *.tiff)"
        )
        if not file_path:
            return

        self.current_image_path = file_path
        self.original_pixmap = QPixmap(file_path)
        scaled_pixmap = self.original_pixmap.scaled(
            self.image_label.size(),
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        )
        self.image_label.setPixmap(scaled_pixmap)
        self.detect_btn.setEnabled(True)
        self.toggle_view_btn.setEnabled(False)  # 检测前禁用切换按钮
        self.results_text.clear()
        self.detection_results = None
        self.show_annotated = True  # 重置显示模式

    def toggle_view(self):
        """切换原图和检测图的显示"""
        if not self.detection_results or not self.original_pixmap:
            return
            
        self.show_annotated = not self.show_annotated
        
        if self.show_annotated:
            # 显示检测图
            annotated_img = self.detection_results["annotated_image"]
            h, w, ch = annotated_img.shape
            q_img = QImage(annotated_img.data, w, h, ch * w, QImage.Format.Format_BGR888)
            pixmap = QPixmap.fromImage(q_img).scaled(
                self.image_label.size(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.image_label.setPixmap(pixmap)
        else:
            # 显示原图
            scaled_pixmap = self.original_pixmap.scaled(
                self.image_label.size(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.image_label.setPixmap(scaled_pixmap)

    def reset_detection_ui(self):
        """重置检测相关UI状态"""
        self.detect_btn.setEnabled(True)
        self.detect_btn.setText("病灶检测")
        self.progress_bar.setVisible(False)

    def show_detection_results(self, result):
        """显示检测结果"""
        self.detection_results = result
        
        # 显示检测图
        annotated_img = result["annotated_image"]
        h, w, ch = annotated_img.shape
        q_img = QImage(annotated_img.data, w, h, ch * w, QImage.Format.Format_BGR888)
        pixmap = QPixmap.fromImage(q_img).scaled(
            self.image_label.size(),
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        )
        self.image_label.setPixmap(pixmap)
        
        # 启用切换视图按钮
        self.toggle_view_btn.setEnabled(True)
        
        # 设置显示模式为检测图
        self.show_annotated = True

        # 显示结果文本
        self.results_text.setText(f"检测到的病灶数量: {result['num_detections']}")